#!/usr/bin/python

import os,numpy,scipy,random,fcntl
from pdb import pm
from multiprocessing import Pool
from pylab import *
from collections import *
from scipy import fftpack
from scipy.cluster import vq
from scipy.ndimage import morphology,filters,interpolation,measurements
from scipy.spatial import distance
from ocrolib import docproc,improc
import pyflann,tables

numpy.seterr(over='raise')
numpy.seterr(divide='raise')
numpy.seterr(invalid='raise')

def make2d(data):
    return data.reshape(len(data),prod(data.shape[1:]))

def showgrid(data,d=None,g=32):
    """Show a data array as a grid of image."""
    clf()
    if d is None: d = int(sqrt(prod(data[0].shape)))
    image = zeros((g*d,g*d))
    for i in range(min(len(data),g*g)):
        im = data[i]
        if im.ndim==1:
            im = im.reshape(d,len(im)/d)
        image[(i/g)*d:(i/g+1)*d,(i%g)*d:(i%g+1)*d] = im
    imshow(image.reshape(g*d,g*d))

def vectorsLessThan(d):
    r = int(d+1)
    result = []
    for i in range(-r,r+1):
        for j in range(-r,r+1):
            if hypot(i,j)>=d: continue
            result.append((i,j))
    return sorted(result,key=lambda x:hypot(*x))

class Shifts:
    def __init__(self,d=1.1):
        self.shifts = vectorsLessThan(d)
    def __len__(self):
        return len(self.shifts)
    def transform(self,image,i):
        delta = self.shifts[i]
        return interpolation.shift(image,delta,order=1)
    
class Skews:
    def __init__(self,n=5,lo=-0.2,hi=0.2):
        self.skews = sorted(linspace(lo,hi,n),key=abs)
    def __len__(self):
        return len(self.skews)
    def transform(self,image,i):
        s = self.skews[i]
        m = array([[1.0,0],[-s,1.0]])
        w,h = image.shape
        c = array([w/2.0,h/2])
        d = c-dot(m,c)
        return interpolation.affine_transform(image,m,offset=d,order=1)

class Aspects:
    def __init__(self,n=3,lo=-0.1,hi=0.1):
        self.aspects = sorted(linspace(lo,hi,n),key=abs)
    def __len__(self):
        return len(self.aspects)
    def transform(self,image,i):
        z = self.aspects[i]
        m = array([[1.0+z,0],[0.0,1.0-z]])
        w,h = image.shape
        c = array([w/2.0,h/2])
        d = c-dot(m,c)
        return interpolation.affine_transform(image,m,offset=d,order=1)

class Zooms:
    def __init__(self,n=5,lo=-0.1,hi=0.3):
        self.zooms = sorted(linspace(lo,hi,n),key=abs)
    def __len__(self):
        return len(self.zooms)
    def transform(self,image,i):
        z = self.zooms[i]
        m = array([[1.0+z,0],[0.0,1.0+z]])
        w,h = image.shape
        c = array([w/2.0,h/2])
        d = c-dot(m,c)
        return interpolation.affine_transform(image,m,offset=d,order=1)

class Dilations:
    def __init__(self,lo=-2,hi=2):
        self.ds = sorted(arange(lo,hi+1),key=abs)
    def __len__(self):
        return len(self.ds)
    def transform(self,image,i):
        d = self.ds[i]
        if d<0:
            return morphology.grey_erosion(image,size=(-d+1,-d+1))
        elif d>0:
            return morphology.grey_dilation(image,size=(d+1,d+1))
        else:
            return image
    
from scipy import stats

def random_transform(image,eps=0.03,delta=0.3,rho=0.1,sigma=1.5):
    m = array([[1+eps*randn(),eps*randn()],[eps*randn(),1.0+eps*randn()]])
    w,h = image.shape
    c = array([w/2.0,h/2])
    d = c-dot(m,c)+array([randn()*delta,randn()*delta])
    image = interpolation.affine_transform(image,m,offset=d,order=1)
    sigma = 0.5+rand()*sigma
    smoothed = filters.gaussian_filter(image,(sigma,sigma))
    return array(smoothed>0.5+(randn()*rho),'f')

def random_variants(image,n=100,**kw):
    image = image*1.0/amax(image)
    result = [image]
    for i in range(3*n):
        distorted = random_transform(image,**kw)
        if amax(distorted)<0.1: continue
        distorted = distorted*1.0/amax(distorted)
        if abs(sum(distorted)-sum(image))/sum(image) > 0.2: continue
        result.append(distorted)
        if len(result)>=n: break
    return result

transforms = [Shifts(),Aspects(),Skews(),Zooms(),Dilations()]
numbers = [len(x) for x in transforms]

def encounts(l):
    if len(l)==1:
        for i in range(l[0]):
            yield [i]
    else:
        for i in range(l[0]):
            for s in encounts(l[1:]):
                yield [i]+s

def variants(image):
    for params in sorted(encounts(numbers),key=norm):
        x = image
        for t,p in zip(transforms,params):
            x = t.transform(x,p)
        yield x/maximum(1e-3,amax(x))

def rvariants(image,n=10000):
    s = sum(image)
    for x in variants(image):
        if n<=0: break
        if sum(x)<s/2: continue
        if sum(x[0])>1: continue
        if sum(x[-1])>1: continue
        yield x
        n -= 1

def multisorted(*args):
    order = argsort(args[0])
    return [a[order] for a in args]

class NNClass:
    def __init__(self):
        with tables.openFile("dkmeans200.h5") as db:
            print "# protos",len(db.root.patches)
            self.data = array(db.root.patches)
            self.classes = array(db.root.pclasses)
            self.flann = pyflann.FLANN()
            self.flann.build_index(make2d(self.data))
            self.k = 1
            self.n = 300
            print "# done"
    def classify(self,image):
        vs = array(list(rvariants(image,self.n)))
        result,dists = self.flann.nn_index(make2d(vs),self.k)
        dists,result = multisorted(dists.ravel(),result.ravel())
        classes = [self.classes[i] for i in result]
        return Counter(classes[:5]).most_common(1)[0][0]

def test():
    print "initializing"
    nnclass = NNClass()
    print "done"
    errors = 0
    count = 0
    for i in range(300000,1500000,117):
        cls = db.root.classes[i]
        pred = nnclass.classify(db.root.patches[i])
        if cls in [ord("~"),ord("_")] or cls>255: continue
        if not confusable(chr(pred),chr(cls)): errors += 1
        count += 1
        print i,chr(cls),chr(pred),errors*1.0/count
